#!/usr/bin/env python3

import asyncio
import subprocess
import os
import pty
import sys
import termios
import struct
import fcntl
import argparse
import signal


def set_winsize(master_fd, row, col):
    """设置 PTY 窗口大小"""
    winsize = struct.pack("HHHH", row, col, 0, 0)
    fcntl.ioctl(master_fd, termios.TIOCSWINSZ, winsize)


def get_winsize(master_fd):
    """获取当前 PTY 窗口大小"""
    try:
        winsize = fcntl.ioctl(master_fd, termios.TIOCGWINSZ, b'\x00' * 8)
        row, col = struct.unpack("HHHH", winsize)[:2]
        return row, col
    except:
        return 24, 80


def setup_pty():
    """子进程启动前配置终端属性"""
    try:
        # 获取 slave 端（0 是 stdin）
        attrs = termios.tcgetattr(0)
        # 设置回显和规范模式
        attrs[3] &= termios.ECHO      # 回显on
        attrs[3] &= ~termios.ICANON    # 原始模式
        attrs[3] &= ~termios.ECHOE     # 禁用退格回显
        attrs[3] &= ~termios.ECHOK     # 禁用换行回显
        attrs[6][termios.VMIN] = 1     # 最小读取字符
        attrs[6][termios.VTIME] = 0    # 超时
        termios.tcsetattr(0, termios.TCSANOW, attrs)
    except:
        pass


async def handle_client(reader, writer):
    client_addr = writer.get_extra_info('peername')
    print(f"[+] 已接受来自 {client_addr[0]}:{client_addr[1]} 的连接")

    # 创建 PTY
    master, slave = pty.openpty()

    # 设置初始窗口大小
    set_winsize(master, 24, 80)

    # 设置环境变量，防止 Vim 输出复杂控制序列
    env = os.environ.copy()
    env["TERM"] = "xterm"        # 关键：避免颜色和图形终端序列，TERM=dumb
    env["PATH"] = "/usr/local/bin:/usr/bin:/bin:/usr/sbin:/sbin"  # 确保 PATH 正常
    env["PS1"] = r"\u@\h:\w\$ "  # 可选：设置简单提示符

    try:
        # 创建子进程
        proc = await asyncio.create_subprocess_exec(
            "/bin/bash",
            stdin=slave,
            stdout=slave,
            stderr=slave,
            env=env,
            preexec_fn=setup_pty  # 子进程启动前配置终端
        )
    except Exception as e:
        print(f"[-] 启动 bash 失败: {e}")
        writer.close()
        await writer.wait_closed()
        os.close(master)
        os.close(slave)
        return

    os.close(slave)  # 关闭 slave 端（由子进程使用）

    # 创建两个任务来处理双向数据传输
    tcp_to_pty_task = asyncio.create_task(tcp_to_pty(reader, master))
    pty_to_tcp_task = asyncio.create_task(pty_to_tcp(master, writer))

    try:
        # 等待任一任务完成
        done, pending = await asyncio.wait(
            [tcp_to_pty_task, pty_to_tcp_task],
            return_when=asyncio.FIRST_COMPLETED
        )

        # 取消未完成的任务
        for task in pending:
            task.cancel()
            try:
                await task
            except asyncio.CancelledError:
                pass
    except Exception as e:
        print(f"[-] 处理客户端连接时出错: {e}")
    finally:
        # 清理资源
        print(f"[-] 连接结束：{client_addr[0]}:{client_addr[1]}")
        writer.close()
        await writer.wait_closed()
        os.close(master)
        
        # 终止子进程
        try:
            proc.terminate()
            await asyncio.wait_for(proc.wait(), timeout=2)
        except asyncio.TimeoutError:
            proc.kill()  # 强制结束进程
        except ProcessLookupError:
            pass  # 进程已经结束


async def tcp_to_pty(reader, master_fd):
    """从TCP连接读取数据并写入PTY"""
    try:
        while True:
            data = await reader.read(4096)
            if not data:
                break
            os.write(master_fd, data)
    except Exception as e:
        pass  # 连接关闭或其他异常


async def pty_to_tcp(master_fd, writer):
    """从PTY读取数据并写入TCP连接"""
    loop = asyncio.get_event_loop()
    try:
        while True:
            # 使用异步方式读取PTY数据
            data = await loop.run_in_executor(None, os.read, master_fd, 4096)
            if not data:
                break
            writer.write(data)
            await writer.drain()
    except Exception as e:
        pass  # 连接关闭或其他异常

async def run(host, port):
    server = await asyncio.start_server(handle_client, host, port)
    addr = server.sockets[0].getsockname()
    print(f"[*] start tiny telnetd on {addr[0]}:{addr[1]} ...")
    async with server:
        await server.serve_forever()

async def main():
    parser = argparse.ArgumentParser(description="Async TCP to Interactive Shell with Terminal Fix")
    parser.add_argument("-b", "--bind", default="0.0.0.0", help="监听地址 (默认: 0.0.0.0)")
    parser.add_argument("-p", "--port", type=int, default=22222, help="监听端口 (默认: 22222)")
    parser.add_argument("--cols", type=int, default=80, help="初始列数 (默认: 80)")
    parser.add_argument("--rows", type=int, default=24, help="初始行数 (默认: 24)")
    args = parser.parse_args()

    try:
        # 启动服务器
        server = await asyncio.start_server(
            handle_client,
            args.bind,
            args.port
        )

        addr = server.sockets[0].getsockname()
        print(f"[*] 正在监听 {addr[0]}:{addr[1]} ...")
        print(f"[*] 使用: nc {addr[0]} {addr[1]} 或 socat -,raw,echo=0 tcp:{addr[0]}:{addr[1]}")

        async with server:
            await server.serve_forever()
    except KeyboardInterrupt:
        print("\n[!] 服务已停止")
    except Exception as e:
        print(f"[-] 错误: {e}")
        sys.exit(1)


if __name__ == "__main__":
    asyncio.run(main())
